#!/usr/bin/python -u
# -*- coding: utf-8 -*-
#
# Copyright 2015 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import print_function

import jsonrpclib
import logging
import os
import signal
import sys

import factory_common  # pylint: disable=W0611
from cros.factory.utils import jsonrpc_utils
from cros.factory.utils import net_utils
from cros.factory.utils import process_utils


class DHCPManager(object):
  """The manager that provides DHCP service.

  Properties set from __init__ arguments:
    interface: The name of the network interface for DHCP.
    my_ip: The IP address for the DHCP server. The interface used will be
      set to this IP address.
    ip_start: The start of DHCP IP range.
    ip_end: The end of DHCP IP range.
    lease_time: How long in seconds before the DHCP lease expires.
    bootp: A tuple (ip, filename, hostname) specifying the boot parameters.
    on_add: Optional callback function that's called when a client is issued
      a new IP address. The callback takes the IP address as the only argument.
    on_old: Similar to on_add, but called when a client is fed an IP address
      that's previously assigned to it.
    on_del: Similar to on_del, but called when a lease is revoked.
  """

  # Where run files are placed
  VARRUN_DIR = '/var/run/dhcp_manager'

  # Callback file name prefix
  CALLBACK_PREFIX = 'dhcp_manager_cb_'

  # Dnsmasq PID file name prefix
  PID_PREFIX = 'dnsmasq_dhcp_manager_pid_'

  # Lease file name prefix
  LEASE_PREFIX = 'dnsmasq_leases_'

  def __init__(self, interface,
               my_ip='192.168.0.1',
               netmask='255.255.255.0',
               ip_start='192.168.0.10',
               ip_end='192.168.0.20',
               lease_time=3600,
               bootp=None,
               on_add=None,
               on_old=None,
               on_del=None):
    self._interface = interface
    self._my_ip = my_ip
    self._netmask = netmask
    self._ip_start = ip_start
    self._ip_end = ip_end
    self._lease_time = lease_time
    self._bootp = bootp
    self._rpc_server = None
    self._process = None
    self._dhcp_action = {'add': on_add, 'old': on_old, 'del': on_del}
    self._callback_port = None

  def StartDHCP(self):
    """Starts DHCP service."""
    if not os.path.exists(self.VARRUN_DIR):
      os.makedirs(self.VARRUN_DIR)
    self._callback_port = net_utils.FindUnusedTCPPort()
    self._rpc_server = jsonrpc_utils.JSONRPCServer(
        port=self._callback_port, methods={'Callback': self.DHCPCallback})
    self._rpc_server.Start()
    # __file__ may be a generated .pyc file that's not executable. Use .py.
    callback_file_target = __file__.replace('.pyc', '.py')
    callback_file_symlink = os.path.join(self.VARRUN_DIR,
                                         '%s%d' % (self.CALLBACK_PREFIX,
                                                   self._callback_port))
    os.symlink(callback_file_target, callback_file_symlink)
    dhcp_range = '%s,%s,%d' % (self._ip_start, self._ip_end, self._lease_time)
    dns_port = net_utils.FindUnusedTCPPort()
    lease_file = os.path.join(self.VARRUN_DIR,
                              '%s%s' % (self.LEASE_PREFIX, self._interface))
    pid_file = os.path.join(self.VARRUN_DIR,
                            '%s%s' % (self.PID_PREFIX, self._interface))
    # Make sure the interface is up
    net_utils.SetEthernetIp(self._my_ip, self._interface, self._netmask)
    # Start dnsmasq and have it call back to us on any DHCP event.
    cmd = ['dnsmasq',
           '--no-daemon',
           '--dhcp-range', dhcp_range,
           '--interface', self._interface,
           '--port', str(dns_port),
           '--no-dhcp-interface=%s' % net_utils.GetDefaultGatewayInterface(),
           '--dhcp-leasefile=%s' % lease_file,
           '--pid-file=%s' % pid_file,
           '--dhcp-script', callback_file_symlink]
    if self._bootp:
      cmd.append('--dhcp-boot=%s,%s,%s' %
                 (self._bootp[1], self._bootp[2], self._bootp[0]))
    self._process = process_utils.Spawn(cmd, sudo=True, log=True)
    # Make sure the IP address is set on the interface
    net_utils.SetEthernetIp(self._my_ip, self._interface, self._netmask,
                            force=True)
    # Make sure DHCP packets are not blocked
    process_utils.Spawn(['iptables',
                         '--insert', 'INPUT',
                         '--protocol', 'udp',
                         '--dport', '67:68',
                         '--sport', '67:68',
                         '-j', 'ACCEPT'],
                        sudo=True)

  def StopDHCP(self):
    """Stops DHCP service."""
    self._process.terminate()
    self._process = None
    self._rpc_server.Destroy()
    self._rpc_server = None
    net_utils.Ifconfig(self._interface, enable=False)
    callback_file_symlink = os.path.join(self.VARRUN_DIR,
                                         '%s%d' % (self.CALLBACK_PREFIX,
                                                   self._callback_port))
    os.unlink(callback_file_symlink)

  def DHCPCallback(self, argv):
    """RPC method that processes the parameters from dnsmasq.

    Based on the parameters, the appropriate callback action is invoked.

    Args:
      argv: The complete arguments received from dnsmasq callback.
    """
    if len(argv) < 4 or argv[1] not in self._dhcp_action:
      logging.error("Invalid DHCP callback: %s", argv)
      return
    action = self._dhcp_action[argv[1]]
    ip = argv[3]
    if action:
      action(ip)

  @classmethod
  def CleanupStaleInstance(cls):
    """Kills all running dnsmasq instance and clean up run directory."""
    if not os.path.exists(cls.VARRUN_DIR):
      return
    for run_file in os.listdir(cls.VARRUN_DIR):
      if run_file.startswith(cls.PID_PREFIX):
        intf = run_file[len(cls.PID_PREFIX):]
        full_run_path = os.path.join(cls.VARRUN_DIR, run_file)
        with open(full_run_path, 'r') as f:
          pid = int(f.read())
        try:
          os.kill(pid, signal.SIGKILL)
        except OSError:
          pass
        net_utils.Ifconfig(intf, enable=False)
        os.unlink(full_run_path)
      if run_file.startswith(cls.CALLBACK_PREFIX):
        os.unlink(os.path.join(cls.VARRUN_DIR, run_file))


class DummyDHCPManager(object):
  """A dummy DHCPManager.

  Replace DHCPManager with this for unit tests so that unit tests don't mess
  with the system network setting.
  """
  def __init__(self, *args, **kwargs):
    pass

  def StartDHCP(self):
    pass

  def StopDHCP(self):
    pass

  @classmethod
  def CleanupStaleInstance(cls):
    pass


if __name__ == '__main__':
  # Figure out what port to call back to
  filename = os.path.basename(sys.argv[0])
  if not filename.startswith(DHCPManager.CALLBACK_PREFIX):
    # Not callback. Do nothing.
    sys.exit(0)
  callback_port = int(filename[len(DHCPManager.CALLBACK_PREFIX):])

  # Forward whatever dnsmasq tells us to DHCPManager
  proxy = jsonrpclib.Server('http://127.0.0.1:%d' % callback_port,
                            transport=jsonrpc_utils.TimeoutJSONRPCTransport(1))
  proxy.Callback(sys.argv)
